# code2/evaluation_utils.py
from collections import Counter
import torch

def get_ngrams(tokens, n):
    return [tuple(tokens[i:i + n]) for i in range(len(tokens) - n + 1)]

def ngram_overlap(a_tokens, b_tokens, n):
    a_ngrams = get_ngrams(a_tokens, n)
    b_ngrams = set(get_ngrams(b_tokens, n))
    if not a_ngrams:
        return 0.0
    overlap = sum(1 for ng in a_ngrams if ng in b_ngrams)
    return overlap / len(a_ngrams)

def compute_ELn(pred_tokens, target_tokens, n=10):
    T = len(target_tokens)
    if T <= n:
        return 0.0

    total_overlap = 0.0
    for t in range(1, T - n + 1):  # 注意 t 从 1 到 T - n
        pred_subseq = pred_tokens[:t + n]  # 模拟基于 x< t 的生成
        target_subseq = target_tokens[t:]
        if len(pred_subseq) < n or len(target_subseq) < n:
            continue
        total_overlap += ngram_overlap(pred_subseq, target_subseq, n)

    return total_overlap / (T - n)

def compute_MA(pred_logits, target_tokens):
    pred_ids = torch.argmax(pred_logits, dim=-1)
    correct = (pred_ids == target_tokens).float()
    if correct[1:].numel() == 0:
        print("[DEBUG] Empty target slice")
        return 0.0
    return correct[1:].mean().item()

